{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyspark import SparkContext, SparkConf\n",
    "from pyspark.sql import SparkSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sc = SparkContext(conf=SparkConf())\n",
    "spark = SparkSession(sparkContext=sc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+------+---+---+---+---+\n",
      "| x1|    x2| x3| x4| y1| y2|\n",
      "+---+------+---+---+---+---+\n",
      "|  a| apple|  1|2.4|  1|yes|\n",
      "|  a|orange|  1|2.5|  0| no|\n",
      "|  b|orange|  2|3.5|  1| no|\n",
      "|  b|orange|  2|1.4|  0|yes|\n",
      "|  b| peach|  2|2.1|  0|yes|\n",
      "|  c| peach|  4|1.5|  1|yes|\n",
      "+---+------+---+---+---+---+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "pdf = pd.DataFrame({\n",
    "        'x1': ['a','a','b','b', 'b', 'c'],\n",
    "        'x2': ['apple', 'orange', 'orange','orange', 'peach', 'peach'],\n",
    "        'x3': [1, 1, 2, 2, 2, 4],\n",
    "        'x4': [2.4, 2.5, 3.5, 1.4, 2.1,1.5],\n",
    "        'y1': [1, 0, 1, 0, 0, 1],\n",
    "        'y2': ['yes', 'no', 'no', 'yes', 'yes', 'yes']\n",
    "    })\n",
    "df = spark.createDataFrame(pdf)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VectorAssembler\n",
    "\n",
    "To fit a ML model in pyspark, we need to combine all feature columns into one single column of vectors: the **featuresCol**. The `VectorAssembler` can be used to combine multiple **`OneHotEncoder` columns** and **other continuous variable columns** into one single column.\n",
    "\n",
    "The example below shows how to combine three OneHotEncoder columns and one numeric column into a **featureCol** column."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## StringIndex and OneHotEncode categorical columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler\n",
    "from pyspark.ml import Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[StringIndexer_4432aeb7f8b4c9e5fad5,\n",
       " StringIndexer_40329fc3347ca61a760a,\n",
       " StringIndexer_470783ab8b97d80f06b0,\n",
       " OneHotEncoder_44af91c8eb1417ffd215,\n",
       " OneHotEncoder_43319000e21f77b48397,\n",
       " OneHotEncoder_41de9c9da21d8f8985a3]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_stages = [StringIndexer(inputCol=c, outputCol='idx_' + c) for c in ['x1', 'x2', 'x3']] + \\\n",
    "             [OneHotEncoder(inputCol='idx_' + c, outputCol='ohe_' + c) for c in ['x1', 'x2', 'x3']]\n",
    "all_stages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+------+---+---+---+---+------+------+------+-------------+-------------+-------------+\n",
      "| x1|    x2| x3| x4| y1| y2|idx_x1|idx_x2|idx_x3|       ohe_x1|       ohe_x2|       ohe_x3|\n",
      "+---+------+---+---+---+---+------+------+------+-------------+-------------+-------------+\n",
      "|  a| apple|  1|2.4|  1|yes|   1.0|   2.0|   1.0|(2,[1],[1.0])|    (2,[],[])|(2,[1],[1.0])|\n",
      "|  a|orange|  1|2.5|  0| no|   1.0|   0.0|   1.0|(2,[1],[1.0])|(2,[0],[1.0])|(2,[1],[1.0])|\n",
      "|  b|orange|  2|3.5|  1| no|   0.0|   0.0|   0.0|(2,[0],[1.0])|(2,[0],[1.0])|(2,[0],[1.0])|\n",
      "|  b|orange|  2|1.4|  0|yes|   0.0|   0.0|   0.0|(2,[0],[1.0])|(2,[0],[1.0])|(2,[0],[1.0])|\n",
      "|  b| peach|  2|2.1|  0|yes|   0.0|   1.0|   0.0|(2,[0],[1.0])|(2,[1],[1.0])|(2,[0],[1.0])|\n",
      "|  c| peach|  4|1.5|  1|yes|   2.0|   1.0|   2.0|    (2,[],[])|(2,[1],[1.0])|    (2,[],[])|\n",
      "+---+------+---+---+---+---+------+------+------+-------------+-------------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_new = Pipeline(stages=all_stages).fit(df).transform(df)\n",
    "df_new.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assemble feature columns into one single **feacturesCol** with **`VectorAssembler`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+------+---+---+---+---+-------------+-------------+-------------+-----------------------------+\n",
      "|x1 |x2    |x3 |x4 |y1 |y2 |ohe_x1       |ohe_x2       |ohe_x3       |featuresCol                  |\n",
      "+---+------+---+---+---+---+-------------+-------------+-------------+-----------------------------+\n",
      "|a  |apple |1  |2.4|1  |yes|(2,[1],[1.0])|(2,[],[])    |(2,[1],[1.0])|(7,[1,5,6],[1.0,1.0,2.4])    |\n",
      "|a  |orange|1  |2.5|0  |no |(2,[1],[1.0])|(2,[0],[1.0])|(2,[1],[1.0])|[0.0,1.0,1.0,0.0,0.0,1.0,2.5]|\n",
      "|b  |orange|2  |3.5|1  |no |(2,[0],[1.0])|(2,[0],[1.0])|(2,[0],[1.0])|[1.0,0.0,1.0,0.0,1.0,0.0,3.5]|\n",
      "|b  |orange|2  |1.4|0  |yes|(2,[0],[1.0])|(2,[0],[1.0])|(2,[0],[1.0])|[1.0,0.0,1.0,0.0,1.0,0.0,1.4]|\n",
      "|b  |peach |2  |2.1|0  |yes|(2,[0],[1.0])|(2,[1],[1.0])|(2,[0],[1.0])|[1.0,0.0,0.0,1.0,1.0,0.0,2.1]|\n",
      "|c  |peach |4  |1.5|1  |yes|(2,[],[])    |(2,[1],[1.0])|(2,[],[])    |(7,[3,6],[1.0,1.5])          |\n",
      "+---+------+---+---+---+---+-------------+-------------+-------------+-----------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_assembled = VectorAssembler(inputCols=['ohe_x1', 'ohe_x2', 'ohe_x3', 'x4'], outputCol='featuresCol')\\\n",
    "    .transform(df_new)\\\n",
    "    .drop('idx_x1', 'idx_x2', 'idx_x3')\n",
    "df_assembled.show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert sparse vectors in featuresCol to dense vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.ml.linalg import SparseVector, DenseVector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def dense_features_col(x):\n",
    "    return(x.toArray().dtype)\n",
    "dense_features_col_udf = udf(dense_features_col, returnType=StringType())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SparseVector(7, {1: 1.0, 5: 1.0, 6: 2.4}),\n",
       " DenseVector([0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 2.5]),\n",
       " DenseVector([1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 3.5]),\n",
       " DenseVector([1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.4])]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_assembled.rdd.map(lambda x: x['featuresCol']).take(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 2.3999999999999999],\n",
       " [0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 2.5],\n",
       " [1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 3.5],\n",
       " [1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.3999999999999999],\n",
       " [1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 2.1000000000000001]]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_assembled.rdd.map(lambda x: list(x['featuresCol'].toArray())).take(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
